import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

# 加载数据集
mnist = input_data.read_data_sets('MNIST_data', one_hot=True)

# 每次执行的个数
batch_size = 100

# 总训练数
train_iters = 100000

# 学习率
lr = 1e-3

# 避免过拟合
drop_out = 0.75

# 定义参数 输入x
x = tf.placeholder(tf.float32, [None, 784])

# 定义参数，正确输出y
y = tf.placeholder(tf.float32, [None, 10])

# 定义参数，drop_out的值
keep_prop = tf.placeholder(tf.float32)

# 定义weight
weights = {

    'w1': tf.Variable(tf.truncated_normal([5, 5, 1, 32])),

    'w2': tf.Variable(tf.truncated_normal([5, 5, 32, 64])),

    'f1': tf.Variable(tf.truncated_normal([7 * 7 * 64, 1024])),

    'f2': tf.Variable(tf.truncated_normal([1024, 10]))

}

# 定义biases
biases = {

    'b1': tf.Variable(tf.constant(0.1, shape=[32])),

    'b2': tf.Variable(tf.constant(0.1, shape=[64])),

    'b3': tf.Variable(tf.constant(0.1, shape=[1024])),

    'b4': tf.Variable(tf.constant(0.1, shape=[10]))

}


# 卷积函数
def conv2d(x, w):
    return tf.nn.conv2d(x, w, strides=[1, 1, 1, 1], padding='SAME')


# 池化操作
def max_pool_2x2(x, k=2):
    return tf.nn.max_pool(x, ksize=[1, k, k, 1], strides=[1, k, k, 1], padding='SAME')


# 卷积神经网络
def CNN(x, weights, biases, keep_prop):
    x = tf.reshape(x, shape=[-1, 28, 28, 1])

    conv1 = tf.nn.relu(conv2d(x, weights['w1']) + biases['b1'])
    pool1 = max_pool_2x2(conv1, k=2)

    conv2 = tf.nn.relu(conv2d(pool1, weights['w2']) + biases['b2'])
    pool2 = max_pool_2x2(conv2, k=2)

    pool2_flat = tf.reshape(pool2, shape=[-1, 7 * 7 * 64])
    f1 = tf.nn.relu(tf.matmul(pool2_flat, weights['f1']) + biases['b3'])

    f1 = tf.nn.dropout(f1, keep_prop)

    out = tf.matmul(f1, weights['f2']) + biases['b4']

    return out


# 预测的值
pred = CNN(x, weights, biases, keep_prop)

# 损失函数
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=pred, labels=y))

# 训练函数
optimizer = tf.train.AdamOptimizer(lr).minimize(cost)

# 预测正确
correct_pred = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1))

# 正确率
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

# 初始化
init = tf.global_variables_initializer()

config = tf.ConfigProto()
config.gpu_options.allow_growth = True

saver = tf.train.Saver()

with tf.Session() as sess:
    print("加载模型")
    saver.restore(sess, "./model/model.ckpt")

    step = 0
    while step * batch_size < train_iters:
        print("step :" + str(step))
        train_x, train_y = mnist.train.next_batch(1)
        result = sess.run(pred, feed_dict={x: train_x, y: train_y, keep_prop: drop_out})
        print("正确值", tf.argmax(train_y, 1).eval())
        print("猜测值", tf.argmax(result, 1).eval())
        print("是否正确", tf.equal(tf.argmax(train_y, 1), tf.argmax(result, 1)).eval())
        step = step + 1
